import os
from uuid import uuid4
from langchain_community.embeddings import QianfanEmbeddingsEndpoint
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import Docx2txtLoader
from langchain_community.vectorstores import Chroma
import streamlit as st
from openai import OpenAI
from langchain.chains import ConversationalRetrievalChain
from langchain_deepseek import ChatDeepSeek
from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate

os.environ["QIANFAN_AK"] = "SGbbQdjFjlKurTfUIjYM0Q4P"
os.environ["QIANFAN_SK"] = "lb1tKvDGRhqLZYH4ZYpke6Vco9n9X8Xv"

llm = ChatDeepSeek(
    model="deepseek-chat",
    api_key="sk-bfdc307c3def4f9da9a06775a127e7a1"
)
# 设置页面标题
st.set_page_config(page_title="AI智能标书", page_icon="📊")

if "messages" not in st.session_state:
    st.session_state.messages = []

with st.sidebar:
    st.markdown("<h1 style='text-align: center;color: red;'>AI智能标书</h1>", unsafe_allow_html=True)

# 2. 再创建输入框（显示在下方）
query = st.chat_input("请输入内容")

# 配置保存路径 1
save_dir = "D:\\hbyt\\AI智能投标\\2025_04_23_Word\\2025_04_23_Word\\Word\\4.0\\中文\\uploaded_file"
os.makedirs(save_dir, exist_ok=True)

# 文件上传组件
uploaded_file = st.file_uploader(
    "上传文件（仅支持docx）",
    type=["docx"]
)


@st.cache_resource
def get_retriver(file_path):
    loader = Docx2txtLoader(file_path)
    pages = loader.load()
    text_spliter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=50,
        length_function=len,
        add_start_index=True
    )
    text_documents = text_spliter.split_documents(pages)
    embeddings = QianfanEmbeddingsEndpoint()
    vectordb = Chroma.from_documents(documents=text_documents, embedding=embeddings,
                                     persist_directory="D:\\hbyt\\project\\aibid\\db\\d")
    retriever = vectordb.as_retriever(search_kwargs={"k": 10})
    return retriever


if uploaded_file is not None:
    # 生成唯一文件名
    original_name, ext = os.path.splitext(uploaded_file.name)
    unique_name = f"{original_name}_{uuid4().hex}{ext}"
    save_path = os.path.join(save_dir, unique_name)
    # 保存文件
    try:
        with open(save_path, "wb") as f:
            f.write(uploaded_file.getbuffer())
        get_retriver(save_path)
    except Exception as e:
        st.error(f"❌ 保存失败：{str(e)}")

if query:
    new_path = "D:\\hbyt\\AI智能投标\\2025_04_23_Word\\2025_04_23_Word\\Word\\1.0\\中文\\EHS相关\\EHS保障体系_v1.0_2211.docx"
    loader = Docx2txtLoader(new_path)
    pages = loader.load()
    text_spliter = RecursiveCharacterTextSplitter(
        chunk_size=384,
        chunk_overlap=50,
        length_function=len,
        add_start_index=True
    )
    text_documents = text_spliter.split_documents(pages)
    embeddings = QianfanEmbeddingsEndpoint()
    vectordb = Chroma.from_documents(documents=text_documents, embedding=embeddings,
                                     persist_directory="D:\\hbyt\\project\\aibid\\db\\d")
    retriever = vectordb.as_retriever(search_kwargs={"k": 10})
    top_chunks = retriever.get_relevant_documents(query)

    for i in range(10):
        path = top_chunks[i].metadata['source']
        # # split分割法（推荐）
        parts = path.split('\\')
        file_name = ''
        if len(parts) == 10:
            file_name = '\\'.join(parts[9:])
        suffix = f"文件名称  : {file_name}"  # 从第4个元素开始拼接
        st.sidebar.header(suffix)
        st.sidebar.markdown(top_chunks[i].page_content)

        windows_path = path  # 注意双反斜杠
        unix_path = windows_path.replace("\\", "/")

        with open(unix_path, "rb") as f:
            file_bytes = f.read()
        # 显示下载按钮
        st.sidebar.download_button(
            label=file_name,
            data=file_bytes,
            file_name=file_name,
            mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
            key = i  # Unique key
        )


    # 构建提示模板
    system_template = """你是一个有帮助的AI智能投标小助手。根据以下上下文回答问题,如果提问的问题不是文档相关问题，禁止回答:
    ----------------
    应该回答如下：
    首先对检索到的每一个片段文档和问题进行最相似的排序，排序完成后，再进行回答，
    文件名称填写到文件名称，如“财务管理_v3.1_2406.docx”，其他总结到文件内容中，
    文件名称第一行，文件内容第二行，回答格式如下：
    
    文件名称：，
    总结如下：
    
    开始回答：

    {context}"""
    messages = [
        SystemMessagePromptTemplate.from_template(system_template),
        HumanMessagePromptTemplate.from_template("{question}")
    ]
    qa_prompt = ChatPromptTemplate.from_messages(messages)
    qa_chain = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=retriever,
        condense_question_llm=llm,
        combine_docs_chain_kwargs={"prompt": qa_prompt},
        return_source_documents=True,
        verbose=True
    )
    # 多轮对话
    chat_history = []
    result = qa_chain({"question": query, "chat_history": chat_history})
    # 更新聊天历史
    chat_history.append((query, result["answer"]))
    st.session_state.messages.append({"role": "user", "content": query})
    st.session_state.messages.append({"role": "assistant", "content": result["answer"]})

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])